import torch
import torchvision
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
import numpy as np
from torch.utils.data import DataLoader
from torchvision.datasets import CIFAR10
from model import ImageClassifier
from utils import load_model
from sklearn.metrics import confusion_matrix
import seaborn as sns

# 定义类别名称
classes = ('plane', 'car', 'bird', 'cat', 'deer', 
           'dog', 'frog', 'horse', 'ship', 'truck')

def visualize_predictions(model, test_loader, device, num_images=9):
    """
    可视化模型预测结果
    
    参数:
        model: 训练好的模型
        test_loader: 测试数据加载器
        device: 设备(CPU/GPU)
        num_images: 要可视化的图像数量
    """
    model.eval()
    images, labels = next(iter(test_loader))
    images, labels = images.to(device), labels.to(device)
    
    # 随机选择一些图片
    indices = torch.randperm(images.size(0))[:num_images]
    images = images[indices]
    labels = labels[indices]

    # 获取预测结果
    with torch.no_grad():
        outputs = model(images)
        _, predicted = torch.max(outputs, 1)
        probabilities = torch.nn.functional.softmax(outputs, dim=1)

    # 显示图片和预测结果
    plt.figure(figsize=(12, 12))
    for i in range(num_images):
        plt.subplot(3, 3, i + 1)
        # 反归一化
        img = images[i].cpu().numpy().transpose((1, 2, 0))
        img = img * np.array([0.2470, 0.2435, 0.2616]) + np.array([0.4914, 0.4822, 0.4465])
        img = np.clip(img, 0, 1)
        
        plt.imshow(img)
        plt.title(f"Pred: {classes[predicted[i]]} ({probabilities[i][predicted[i]]:.2f})\nTrue: {classes[labels[i]]}")
        plt.axis('off')
    plt.tight_layout()
    plt.savefig('predictions.png')
    plt.close()

def plot_confusion_matrix(model, test_loader, device):
    """
    绘制混淆矩阵
    
    参数:
        model: 训练好的模型
        test_loader: 测试数据加载器
        device: 设备(CPU/GPU)
    """
    model.eval()
    all_preds = []
    all_labels = []
    
    with torch.no_grad():
        for images, labels in test_loader:
            images = images.to(device)
            labels = labels.to(device)
            outputs = model(images)
            _, preds = torch.max(outputs, 1)
            
            all_preds.extend(preds.cpu().numpy())
            all_labels.extend(labels.cpu().numpy())
    
    # 计算混淆矩阵
    cm = confusion_matrix(all_labels, all_preds)
    
    # 绘制混淆矩阵
    plt.figure(figsize=(10, 8))
    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', 
                xticklabels=classes, yticklabels=classes)
    plt.xlabel('Predicted')
    plt.ylabel('True')
    plt.title('Confusion Matrix')
    plt.tight_layout()
    plt.savefig('confusion_matrix.png')
    plt.close()

if __name__ == "__main__":
    # 设置设备
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Using device: {device}")

    # 加载模型
    model = ImageClassifier().to(device)
    load_model(model, "improved_cifar10_classifier.pth")
    print("Model loaded successfully")

    # 加载测试数据
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616))
    ])
    test_dataset = CIFAR10(root='./data', train=False, download=True, transform=transform)
    test_loader = DataLoader(test_dataset, batch_size=100, shuffle=True)

    # 可视化预测结果
    visualize_predictions(model, test_loader, device)
    
    # 绘制混淆矩阵
    plot_confusion_matrix(model, test_loader, device)